package com.kira.scaffoldmvc.ShutDownHook;

import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import org.slf4j.MDC;
import java.util.Map;
import java.util.concurrent.*;

public class TtlThreadPoolExecutor extends ThreadPoolExecutor {
    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }


    @Override
    public void execute(Runnable command) {
        // 获取当前线程的MDC上下文
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        
        // 使用TtlRunnable包装，确保上下文传递
        super.execute(TtlRunnable.get(() -> {
            // 恢复MDC上下文
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
            try {
                command.run();
            } finally {
                // 不要在这里清除MDC，可能会影响后续任务
                // 由任务本身决定何时清理
            }
        }));
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(wrapTask(task));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(wrapTask(task));
    }

    private Runnable wrapTask(Runnable task) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return TtlRunnable.get(() -> {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
            try {
                task.run();
            } finally {
                MDC.clear();
            }
        });
    }

    private <T> Callable<T> wrapTask(Callable<T> task) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return TtlCallable.get(() -> {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
            try {
                return task.call();
            } finally {
                MDC.clear();
            }
        });
    }
}